//Author: Austen Higgins-Cassidy
// 4/10/2012
#include "assembler.h"
#include <sstream>


std::string op_tok[] = {
	"JSR",
	"INT",
	"IAG",
	"IAS",
	"RFI",
	"IAQ",
	"HWN",
	"HWQ",
	"HWI",
	"SET",
	"ADD",
	"SUB",
	"MUL",
	"MIL",
	"DIV",
	"DVI",
	"MOD",
	"MDI",
	"AND",
	"BOR",
	"XOR",
	"SHR",
	"ASR",
	"SHL",
	"IFB",
	"IFC",
	"IFE",
	"IFN",
	"IFG",
	"IFA",
	"IFL",
	"IFU",
	"ADX",
	"SBX",
	"STI",
	"STD",
	":",
	";"
};

std::string val_tok[] = {
	"A",
	"B",
	"C",
	"X",
	"Y",
	"Z",
	"I",
	"J",
	"EX",
	"PC",
	"SP",
	"[",
	"PUSH",
	"POP",
	"PICK"
};

std::string close_tok[] = {
	"+",
	"]"
};

unsigned short map_op(std::string &tok){
	if(tok == "SET"){
		return 0x1;
	}else if(tok == "ADD"){
		return 0x2;
	}else if(tok == "SUB"){
		return 0x3;
	}else if(tok == "MUL"){
		return 0x4;
	}else if(tok == "MLI"){
		return 0x5;
	}else if(tok == "DIV"){
		return 0x6;
	}else if(tok == "DVI"){
		return 0x7;
	}else if(tok == "MOD"){
		return 0x8;
	}else if(tok == "MDI"){
		return 0x9;
	}else if(tok == "AND"){
		return 0xA;
	}else if(tok == "BOR"){
		return 0xB;
	}else if(tok == "XOR"){
		return 0xC;
	}else if(tok == "SHR"){
		return 0xD;
	}else if(tok == "ASR"){
		return 0xE;
	}else if(tok == "SHL"){
		return 0xF;
	}else if(tok == "IFB"){
		return 0x10;
	}else if(tok == "IFC"){
		return 0x11;
	}else if(tok == "IFE"){
		return 0x12;
	}else if(tok == "IFN"){
		return 0x13;
	}else if(tok == "IFG"){
		return 0x14;
	}else if(tok == "IFA"){
		return 0x15;
	}else if(tok == "IFL"){
		return 0x16;
	}else if(tok == "IFU"){
		return 0x17;
	}else if(tok == "ADX"){
		return 0x1a;
	}else if(tok == "SBX"){
		return 0x1b;
	}else if(tok == "STI"){
		return 0x1e;
	}else if(tok == "STD"){
		return 0x1f;
	}else if(tok == "JSR"){
		return 0x20;
	}else if(tok == "INT"){
		return 0x27;
	}else if(tok == "IAG"){
		return 0x28;
	}else if(tok == "IAS"){
		return 0x29;
	}else if(tok == "RFI"){
		return 0x2a;
	}else if(tok == "IAQ"){
		return 0x2b;
	}else if(tok == "HWN"){
		return 0x2f;
	}else if(tok == "HWQ"){
		return 0x30;
	}else if(tok == "HWI"){
		return 0x31;
	}else return 0xFF;

}

unsigned short map(std::string &buf, inst *i,  std::vector<std::pair<std::string, unsigned short*>> *lct){
	unsigned short size = 1;
	std::string a = buf.substr(0, buf.find_first_of(','));
	buf = buf.substr(buf.find_first_of(',')+1);
	if(i->op != 0x0000)
		i->op |= ((map_val(a, &(i->b), lct)<<5)&OP_B);
	i->op |= ((map_val(buf, &(i->a), lct)<<10)&OP_A);
	if(i->a != 0)
		++size;
	if(i->b != 0)
		++size;
	return size;
	//The op is now fully encoded and we've returned the # words required
}

unsigned short map_val(std::string tok, unsigned short *word,  std::vector<std::pair<std::string, unsigned short*>> *lct){
	unsigned short res = 0xFF;
	bool ram = false;
	bool neg = false;
	bool wreg = false;
	bool no_comp = false; //Do not attempt to compress the value
	std::string nxt;
	chomp(tok);
	if(tok[0] == '-'){
		neg = true;
	}
	if(tok[0] == '['){
		ram = true;
		tok = tok.substr(1);
	}
	if(tok.find('+') != std::string::npos){
		wreg = true;
		nxt = tok.substr(0, tok.find_first_of('+'));
		tok = tok.substr(tok.find_first_of('+')+1);
	}
	chomp(tok);
	if(tok.length() > 0){
		
		if(tok.substr(0,1) == "A"){
			res = 0x00;
			no_comp = true;
		}else if(tok.substr(0,1) == "B"){
			res = 0x01;
			no_comp = true;
		}else if(tok.substr(0,1) == "C"){
			res = 0x02;
			no_comp = true;
		}else if(tok.substr(0,1) == "X"){
			res = 0x03;
			no_comp = true;
		}else if(tok.substr(0,1) == "Y"){
			res = 0x04;
			no_comp = true;
		}else if(tok.substr(0,1) == "Z"){
			res = 0x05; 
			no_comp = true;
		}else if(tok.substr(0,1) == "I"){
			res = 0x06;
			no_comp = true;
		}else if(tok.substr(0,1) == "J"){
			res = 0x07;
			no_comp = true;
		}else if(tok.length() > 1){
			if(tok.substr(0,2) == "EX"){
			res = 0x1d;
			no_comp = true;
			}else if(tok.substr(0,2) == "SP"){
				res = 0x1b;
				return res;
			}else if(tok.substr(0,2) == "PC"){
				res = 0x1c;
				return res;
			}else if(tok.length() > 2){
				if(tok.substr(0,3) == "POP"){
					res =  0x18;
					return res;
				}else if(tok.length() > 3) {
					if(tok.substr(0,4) == "PEEK"){
						res = 0x19;
						return res;
					}else if(tok.substr(0,4) == "PUSH"){
						res = 0x18;
						return res;
					}else if(tok.substr(0,4) == "PICK"){
						res = 0x1a;
						return res;
					}
				}
			}
		}

	}
	if(res < 0x18){
		if(wreg){
			tok = nxt;
			res += 0x10;
		}else if(ram)
			return res + 0x08;

		(*word) = (shtoi(tok));
		if(neg){
			//flip the sign
			(*word) = -((short)(*word));
		}
		//Check to see if word is between -1 and 30
		if(((*word) <= 30 && (*word) >= -1)){
			//store as a literal if we can!
			
			if(!no_comp)
			{
				
				//short s_value;
				res = (*word)+0x21; //Encode correctly (-1 = 0x20, 30 = 0x1e
				//res = res - 2; //(shift to correct for sign)
				(*word) = 0;
			}
		}
		
		return res;
	}
	//Before we assume a number, look for labels
	if(tok[0] >= '0' && tok[0] <= '9') 
		(*word) = (shtoi(tok));
	else{
		//Make an entry in the loc_table
		if(!no_comp){
			//(*lct)[tok] = word;
			std::pair<std::string, unsigned short*> s;
			s.first = tok;
			s.second = word;
			lct->push_back(s);
			res = 0x1f;
			return res;
		}
		*word = 0xFFFF; //LOCKOUT NONZERO WORD
		no_comp = true; //We don't support short-form literal compression
	}


	if(ram)
		res = 0x1e;
	else if(((*word) <= 0x1f) && !no_comp && res == 0xFF){
		res = (*word)+0x20;
		(*word) = 0;
	}else if(res == 0xFF) res = 0x1f;
	return res;
}

char ToHexChar(char t){
	switch(t){
	case 0:
		return '0';
	case 1:
		return '1';
	case 2:
		return '2';
	case 3:
		return '3';
	case 4:
		return '4';
	case 5:
		return '5';
	case 6:
		return '6';
	case 7:
		return '7';
	case 8:
		return '8';
	case 9:
		return '9';
	case 10:
		return 'a';
	case 11:
		return 'b';
	case 12:
		return 'c';
	case 13:
		return 'd';
	case 14:
		return 'e';
	case 15:
		return 'f';
	default:
		return 'X';
	}
}

void WriteHex(std::ostream &out, const unsigned short &word){
	//Write out a 4-character hex string
	//std::string buf = "0x";
	out << " 0x";
	out << ToHexChar((word & 0xF000)>>12);
	out << ToHexChar((word & 0x0F00)>>8);
	out << ToHexChar((word & 0x00F0)>>4);
	out << ToHexChar((word & 0x000F));
	out << " ";

}


void assemble(std::istream &in, std::ostream &out){
	std::string buffer = "";
	std::string token = ";";
	std::string next_token = "";
	inst *instr;
	unsigned short loc = 0;
	bool done = false;
	char res;
	std::vector<inst*> binary;
	std::map<std::string, unsigned short> labels;
	std::vector<std::pair<std::string, unsigned short*>> loc_tab;
	//Assume we have well formed assembly
	while(!done){
		res = extractUntilTokenList(in, buffer, op_tok, dimsize(op_tok), &next_token);
		if(res == -1)
			done = true;
		instr = new inst;
		instr->op = 0;
		instr->a = 0;
		instr->b = 0;
		
		chomp(next_token);
		if(token[0] == ';'){
			token = next_token;
			buffer = "";
			continue;
		}
		if(token[0] == ':'){
			chomp(buffer);
			labels[buffer] = loc;
			token = next_token;
			buffer = "";
			continue;
		}

		instr->op = map_op(token);
		loc += map(buffer, instr, &loc_tab);
		token = next_token;
		buffer = "";
		binary.push_back(instr);
		if(next_token[0] == ';')
			extractUntil(in, buffer, "\n");
		//if(instr->a)
		//	binary.push_back(instr.a);
		//if (instr.b)
		//	binary.push_back(instr.b);
	}//we now have a binary list;
	//Process the values out to words

	//translate the loc tab
	for(std::vector<std::pair<std::string, unsigned short *>>::iterator iter = loc_tab.begin(); iter != loc_tab.end(); ++iter){
		*(iter->second) = labels[iter->first];
	}

	for(std::vector<inst*>::iterator iter = binary.begin(); iter != binary.end(); ++iter){
		
		WriteHex(std::cout, (*iter)->op);
		out.write((char*)&(*iter)->op, sizeof(unsigned short));
		if((*iter)->a != 0){
			WriteHex(std::cout, (*iter)->a);
			out.write((char*)&(*iter)->a, sizeof(unsigned short));
		}
		if((*iter)->b != 0){
			WriteHex(std::cout, (*iter)->b);
			out.write((char*)&(*iter)->b, sizeof(unsigned short));
		}
		delete (*iter);
	}
}

void assemble(std::istream &in, std::ostream &out, unsigned short *&bin, unsigned short &size){
	std::string buffer = "";
	std::string token = ";";
	std::string next_token = "";
	inst *instr;
	unsigned short loc = 0;
	bool done = false;
	char res;
	std::vector<inst*> binary;
	std::map<std::string, unsigned short> labels;
	std::vector<std::pair<std::string, unsigned short*>> loc_tab;
	//Assume we have well formed assembly
	while(!done){
		res = extractUntilTokenList(in, buffer, op_tok, dimsize(op_tok), &next_token);
		if(res == -1)
			done = true;
		instr = new inst;
		instr->op = 0;
		instr->a = 0;
		instr->b = 0;
		
		chomp(next_token);
		if(token[0] == ';'){
			token = next_token;
			buffer = "";
			continue;
		}
		if(token[0] == ':'){
			chomp(buffer);
			labels[buffer] = loc;
			token = next_token;
			buffer = "";
			continue;
		}

		instr->op = map_op(token);
		loc += map(buffer, instr, &loc_tab);
		token = next_token;
		buffer = "";
		binary.push_back(instr);
		if(next_token[0] == ';')
			extractUntil(in, buffer, "\n");
		//if(instr->a)
		//	binary.push_back(instr.a);
		//if (instr.b)
		//	binary.push_back(instr.b);
	}//we now have a binary list;
	//Process the values out to words

	//translate the loc tab
	for(std::vector<std::pair<std::string, unsigned short *>>::iterator iter = loc_tab.begin(); iter != loc_tab.end(); ++iter){
		*(iter->second) = labels[iter->first];
	}
	size = loc;
	bin = new unsigned short[loc];
	int index = 0;
	for(std::vector<inst*>::iterator iter = binary.begin(); iter != binary.end(); ++iter){
		
		WriteHex(std::cout, (*iter)->op);
		bin[index] = (*iter)->op;
		out.write((char*)&(*iter)->op, sizeof(unsigned short));
		if((*iter)->a != 0){
			WriteHex(std::cout, (*iter)->a);
			out.write((char*)&(*iter)->a, sizeof(unsigned short));
			bin[++index] = (*iter)->a;
		}
		if((*iter)->b != 0){
			WriteHex(std::cout, (*iter)->b);
			out.write((char*)&(*iter)->b, sizeof(unsigned short));
			bin[++index] = (*iter)->b;
		}
		++index;
		delete (*iter);
		
	}
}

void writeBin(unsigned short size, unsigned short *bin, std::ostream &out){
		out.write((char*)bin, sizeof(unsigned short)*size);
}
